{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "1nb_Ppf2ZUQL"
   },
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/AI-Hypercomputer/maxtext/blob/main/src/MaxText/examples/sft_qwen3_demo.ipynb)\n",
    "\n",
    "# Qwen3-0.6B Supervised Fine-Tuning (SFT) Demo\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FGbe4_YQZUQL"
   },
   "source": [
    "## Overview\n",
    "\n",
    "This notebook performs SFT training and evaluation workflow on [OpenAI's GSM8K dataset](https://huggingface.co/datasets/openai/gsm8k).\n",
    "The primary goal is to demonstrate the end-to-end process of:\n",
    "1. Pre-SFT Evaluation: Calcuating baseline accuracy for the model before training.\n",
    "2. SFT Training: Fine-tune the model using MaxText & Tunix SFT trainer.\n",
    "3. Post-SFT Evaluation: Re-running the evaluation loop after training to measure the performance gain achieved by SFT.\n",
    "\n",
    "This notebook can run on the **public TPU v5e-1**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zolxPWhQZUQL"
   },
   "source": [
    "## Prerequisites\n",
    "\n",
    "### Change Runtime Type\n",
    "\n",
    "**Instructions:**\n",
    "1.  Navigate to the menu at the top of the screen.\n",
    "2.  Click on **Runtime**.\n",
    "3.  Select **Change runtime type** from the dropdown menu.\n",
    "4.  Select **v5e-1 TPU** as the **Hardware accelerator**.\n",
    "5. Click on **Save**.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Rk_QpVVuZUQL"
   },
   "source": [
    "### Get Your Hugging Face Token\n",
    "\n",
    "To access model checkpoint from the Hugging Face Hub, you need to authenticate with a personal access token.\n",
    "\n",
    "**Follow these steps to get your token:**\n",
    "\n",
    "1.  **Navigate to the Access Tokens page** in your Hugging Face account settings. You can go there directly by visiting this URL:\n",
    "    *   [https://huggingface.co/settings/tokens](https://huggingface.co/settings/tokens)\n",
    "\n",
    "2.  **Create a new token** by clicking the **\"+ Create new token\"** button.\n",
    "\n",
    "3.  **Give your token a name** and assign it a **`read` role**. The `read` role is sufficient for downloading models.\n",
    "\n",
    "4.  **Copy the generated token**. You will need to paste it in the next step.\n",
    "\n",
    "**Follow these steps to store your token:**\n",
    "\n",
    "1. On the left sidebar of your Colab window, click the key icon (the Secrets tab).\n",
    "\n",
    "2. Click **\"+ Add new secret\"**.\n",
    "\n",
    "3. Set the Name as **HF_TOKEN**.\n",
    "\n",
    "4. Paste your token into the Value field.\n",
    "\n",
    "5. Ensure the Notebook access toggle is turned On."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "D9ms-jTSZUQL"
   },
   "source": [
    "## Installation: MaxText & Other Dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**⚠️ Note:** The installation process in following cell may take a few minutes to complete. Please be patient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OSPRVbi7n6tB"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/AI-Hypercomputer/maxtext.git\n",
    "%cd /content/maxtext\n",
    "\n",
    "# Install uv, a fast Python package installer\n",
    "!pip install uv\n",
    "\n",
    "# Install MaxText and its dependencies\n",
    "!uv pip install -e .[tpu] --resolution=lowest\n",
    "!python3 -m MaxText.install_maxtext_extra_deps\n",
    "\n",
    "# Install vLLM for Jax and TPUs\n",
    "!uv pip install vllm-tpu\n",
    "!uv pip install --no-deps qwix==0.1.4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ywtealAxZUQM"
   },
   "source": [
    "### Restart Session\n",
    "To apply certain changes, you need to restart the session.\n",
    "\n",
    "**Instructions:**\n",
    "1.  Navigate to the menu at the top of the screen.\n",
    "2.  Click on **Runtime**.\n",
    "3.  Select **Restart session** from the dropdown menu.\n",
    "\n",
    "You will be asked to confirm the action in a pop-up dialog. Click on **Yes**."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Clexf-j7ZUQM"
   },
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PkBI9A3JZUQM"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import transformers\n",
    "\n",
    "import MaxText\n",
    "from MaxText import pyconfig\n",
    "from MaxText.examples.sft_train_and_evaluate import evaluate_model, get_test_dataset\n",
    "from MaxText.integration.tunix.tunix_adapter import TunixMaxTextAdapter\n",
    "from MaxText.sft import sft_trainer\n",
    "\n",
    "# Suppress vLLM logging with a severity level below ERROR\n",
    "os.environ[\"VLLM_LOGGING_LEVEL\"] = \"ERROR\"\n",
    "from tunix.rl.rollout.vllm_rollout import VllmRollout\n",
    "\n",
    "from datetime import datetime\n",
    "from flax import nnx\n",
    "from huggingface_hub import login\n",
    "\n",
    "try:\n",
    "  from google.colab import userdata\n",
    "except Exception:\n",
    "  pass\n",
    "\n",
    "MAXTEXT_REPO_ROOT = os.path.dirname(MaxText.__file__)\n",
    "print(f\"MaxText installation path: {MAXTEXT_REPO_ROOT}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JBbPN-uVZUQM"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "  HF_TOKEN = userdata.get(\"HF_TOKEN\")\n",
    "except Exception:\n",
    "  HF_TOKEN = os.environ.get(\"HF_TOKEN\")\n",
    "if HF_TOKEN:\n",
    "  login(token=HF_TOKEN)\n",
    "  print(\"Authenticated with Hugging Face successfully!\")\n",
    "else:\n",
    "  print(\"Authentication failed: HF_TOKEN is not set.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "aENuzm9iZUQM"
   },
   "source": [
    "## Model Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RjPYYl3zZUQM"
   },
   "outputs": [],
   "source": [
    "MODEL_NAME = \"qwen3-0.6b\"\n",
    "TOKENIZER_PATH = \"Qwen/Qwen3-0.6B\"\n",
    "tokenizer = transformers.AutoTokenizer.from_pretrained(\n",
    "    TOKENIZER_PATH,\n",
    "    token=HF_TOKEN,\n",
    ")\n",
    "\n",
    "# This colab will download the checkpoint from HF and store at `MODEL_CHECKPOINT_PATH`\n",
    "MODEL_CHECKPOINT_PATH = f\"{MAXTEXT_REPO_ROOT}/qwen_checkpoint\"\n",
    "\n",
    "RUN_NAME = datetime.now().strftime(\"%Y-%m-%d-%H-%m-%S\")\n",
    "\n",
    "# This is the directory where the fine-tuned model checkpoint will be saved\n",
    "BASE_OUTPUT_DIRECTORY = f\"/tmp/maxtext_qwen06\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4L37Ij4NZUQM"
   },
   "source": [
    "## Download Qwen3-0.6B Model Checkpoint from Hugging Face"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kJanDAc0ZUQM"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "!python3 -m MaxText.utils.ckpt_conversion.to_maxtext \\\n",
    "    $MAXTEXT_REPO_ROOT/configs/base.yml \\\n",
    "    model_name=$MODEL_NAME \\\n",
    "    base_output_directory=$MODEL_CHECKPOINT_PATH \\\n",
    "    hf_access_token=$HF_TOKEN \\\n",
    "    use_multimodal=false \\\n",
    "    scan_layers=true"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1T0bm82lZUQM"
   },
   "outputs": [],
   "source": [
    "print(f\"Model checkpoint can be found at: {MODEL_CHECKPOINT_PATH}/0/items\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PC-hILG0ZUQM"
   },
   "source": [
    "## Dataset Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O3MLdr9kZUQM"
   },
   "outputs": [],
   "source": [
    "DATASET_NAME = \"openai/gsm8k\"\n",
    "TRAIN_DATA_SPLIT = \"train\"\n",
    "TEST_DATA_SPLIT = \"test\"\n",
    "HF_DATA_DIR = \"main\"\n",
    "TRAIN_DATA_COLUMNS = [\"question\", \"answer\"]\n",
    "CHAT_TEMPLATE_PATH = f\"{MAXTEXT_REPO_ROOT}/examples/chat_templates/math_qa.json\"\n",
    "NUM_TEST_SAMPLES = 20  # Total number of samples to test\n",
    "BATCH_SIZE = 1  # Number of test samples to process in a batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yeAHmxSYZUQM"
   },
   "source": [
    "## MaxText Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "In-jdp1AAwrL"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "config = pyconfig.initialize(\n",
    "    [\n",
    "        \"\",\n",
    "        f\"{MAXTEXT_REPO_ROOT}/configs/sft.yml\",\n",
    "        f\"load_parameters_path={MODEL_CHECKPOINT_PATH}/0/items\",\n",
    "        f\"model_name={MODEL_NAME}\",\n",
    "        f\"hf_access_token={HF_TOKEN}\",\n",
    "        f\"base_output_directory={BASE_OUTPUT_DIRECTORY}\",\n",
    "        f\"run_name={RUN_NAME}\",\n",
    "        f\"tokenizer_path={TOKENIZER_PATH}\",\n",
    "        f\"hf_path={DATASET_NAME}\",\n",
    "        f\"train_split={TRAIN_DATA_SPLIT}\",\n",
    "        f\"hf_data_dir={HF_DATA_DIR}\",\n",
    "        f\"train_data_columns={TRAIN_DATA_COLUMNS}\",\n",
    "        \"steps=500\",\n",
    "        \"per_device_batch_size=1\",\n",
    "        \"max_target_length=1024\",\n",
    "        \"learning_rate=3e-6\",\n",
    "        \"weight_dtype=bfloat16\",\n",
    "        \"dtype=bfloat16\",\n",
    "        f\"chat_template_path={CHAT_TEMPLATE_PATH}\",\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "O9b0GWo-ZUQM"
   },
   "source": [
    "## Initial Setup & Data Preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TDqFmvUCZUQM"
   },
   "source": [
    "### Create Test Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wscWYxrtZUQM"
   },
   "outputs": [],
   "source": [
    "test_dataset = get_test_dataset(config, tokenizer)\n",
    "test_dataset = test_dataset[:NUM_TEST_SAMPLES]\n",
    "test_dataset = test_dataset.to_iter_dataset().batch(BATCH_SIZE, drop_remainder=True)\n",
    "TOTAL_BATCHES = NUM_TEST_SAMPLES // BATCH_SIZE\n",
    "print(\n",
    "    f\"Processing {NUM_TEST_SAMPLES} examples with a batch size of {BATCH_SIZE}. This will result in {TOTAL_BATCHES} total batches for the test run.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bLSvOOEUZUQM"
   },
   "source": [
    "### Create SFT Trainer State"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2IHsC0m6ZUQM"
   },
   "outputs": [],
   "source": [
    "trainer, mesh = sft_trainer.setup_trainer_state(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PpKtEqzFZUQM"
   },
   "source": [
    "### Create vLLM Rollout"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "3-pf_rbqZUQM"
   },
   "outputs": [],
   "source": [
    "tunix_model = TunixMaxTextAdapter(trainer.model)\n",
    "vllm_rollout = VllmRollout(\n",
    "    model=tunix_model,\n",
    "    tokenizer=tokenizer,\n",
    "    cache_config_or_size=1280,\n",
    "    mesh=mesh,\n",
    "    model_version=TOKENIZER_PATH,\n",
    "    hbm_utilization=0.8,\n",
    "    init_with_random_weights=True,\n",
    "    tpu_backend_type=\"jax\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "567gTxsEZUQM"
   },
   "source": [
    "## Evaluation before SFT Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OnACa3zCZUQM"
   },
   "outputs": [],
   "source": [
    "print(\"Running Pre-SFT Evaluation...\")\n",
    "score = evaluate_model(test_dataset, vllm_rollout, debug=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u5-M4iYkZUQN"
   },
   "outputs": [],
   "source": [
    "print(\"========================= Score for PRE-SFT Evaluation =========================\")\n",
    "print(f\"Percentage of test samples where the model produced the correct numerical answer: {score['correct']}%\")\n",
    "print(\n",
    "    f\"Percentage of test samples where the model produced the numerical answer within 10%: {score['partially_correct']}%\"\n",
    ")\n",
    "print(\n",
    "    f\"Percentage of test samples where the model's output adheres to the expected structure: {score['correct_format']}%\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EJE1ookSAzz-"
   },
   "source": [
    "## SFT Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "id": "mgwpNgQYCJEd",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"Starting SFT Training...\")\n",
    "trainer = sft_trainer.train_model(config, trainer, mesh)\n",
    "print(\"SFT Training Complete!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WEdNYRhwZUQN"
   },
   "source": [
    "## Evaluation after SFT Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XcsZacZdZUQN"
   },
   "outputs": [],
   "source": [
    "print(\"Running Post-SFT Evaluation...\")\n",
    "model = TunixMaxTextAdapter(trainer.model)\n",
    "state = nnx.state(model)\n",
    "vllm_rollout.update_params(state)\n",
    "score = evaluate_model(test_dataset, vllm_rollout, debug=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "editable": true,
    "id": "-JtYTPvJZUQN",
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"========================= Score for POST-SFT Evaluation =========================\")\n",
    "print(f\"Percentage of test samples where the model produced the correct numerical answer: {score['correct']}%\")\n",
    "print(\n",
    "    f\"Percentage of test samples where the model produced the numerical answer within 10%: {score['partially_correct']}%\"\n",
    ")\n",
    "print(\n",
    "    f\"Percentage of test samples where the model's output adheres to the expected structure: {score['correct_format']}%\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "gpuType": "V5E1",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
